package cn.fxc;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.math3.distribution.NormalDistribution;
import org.apache.commons.math3.distribution.TDistribution;
import org.apache.commons.math3.linear.MatrixUtils;
import org.apache.commons.math3.linear.RealMatrix;
import org.apache.commons.math3.stat.correlation.PearsonsCorrelation;
import org.apache.commons.math3.stat.correlation.SpearmansCorrelation;
import org.apache.commons.math3.stat.ranking.NaturalRanking;
import org.apache.commons.math3.stat.ranking.RankingAlgorithm;
import tech.tablesaw.api.Table;
import tech.tablesaw.plotly.api.ScatterPlot;
import tech.tablesaw.plotly.components.Figure;

import java.util.*;

@Slf4j
public class CorrelationCompute {
    RealMatrix correlationMatrix;
    List<Map<String, String>> correlationResultList;

    public static void main(String[] args) {
        Table dataTable = Table.read().csv("data.csv");
        double[][] data = new double[dataTable.columnCount()][];
        for (int i = 0; i < dataTable.columnCount(); i++) {
            data[i] = dataTable.nCol(i).asDoubleArray();
        }
        CorrelationCompute cc= new CorrelationCompute();
//        cc.PearsonCorrelation(data);
        cc.SpearmanCorrelation(data);
        List<Map<String, String>> resultList = cc.correlationResultList;
        TableLogger.logAsTable(resultList);
    }

    public void SpearmanCorrelation(double[][] data){
        correlationResultList = new ArrayList<>();
        RealMatrix dataMatrix = MatrixUtils.createRealMatrix(data).transpose();
        SpearmansCorrelation spearmansCorrelation = new SpearmansCorrelation(dataMatrix);
        correlationMatrix = spearmansCorrelation.getCorrelationMatrix();
        int n =dataMatrix.getRowDimension();
        int m = dataMatrix.getColumnDimension();
        double[][] pValueMatrix = new double[m][m];

        double confidenceLevel = 0.95; // 置信水平
        NormalDistribution normalDistribution = new NormalDistribution(0, 1);
        double criticalValue = normalDistribution.inverseCumulativeProbability(1 - (1 - confidenceLevel) / 2);
        System.out.println("n:"+n);
        System.out.println("criticalValue:"+criticalValue);
        TDistribution tDistribution = new TDistribution(n - 2);

        for (int i = 0; i < m; i++) {
            for (int j = i+1; j < m; j++) {
                if (i == j) {
                    pValueMatrix[i][j] = 0.0; // 对角线上p值为0
                } else {
                    double rho = correlationMatrix.getEntry(i,j);
                    if (Double.isNaN(rho)) {
                        pValueMatrix[i][j] = Double.NaN;
                    } else {
                        double t = rho * Math.sqrt((n - 2) / (1 - rho * rho));
                        pValueMatrix[i][j] = 2 * tDistribution.cumulativeProbability(-Math.abs(t));
//                        double r = correlationMatrix.getEntry(i,j);
                        double[] ci = calculateConfidenceIntervalOfSpearman(rho, n, criticalValue);
                        log.debug("Correlation between {} and {}: {}, 95% CI: [{}, {}], pValue:{}\n", i, j, rho, ci[0], ci[1],pValueMatrix[i][j]);

                        Map<String,String> correlationResult = new LinkedHashMap<>();
                        correlationResult.put("column 1", String.valueOf(i));
                        correlationResult.put("column 2", String.valueOf(j));
                        correlationResult.put("N", String.valueOf(n));
                        correlationResult.put("r", String.format("%.4f",rho));
                        correlationResult.put("ci", "["+String.format("%.4f",ci[0])+","+String.format("%.4f",ci[1])+"]");
                        correlationResult.put("pValue", String.format("%.4f",pValueMatrix[i][j]));
                        correlationResultList.add(correlationResult);
                    }
                }
            }
        }

    }


    public void PearsonCorrelation(double[][] data){
        RealMatrix dataMatrix = MatrixUtils.createRealMatrix(data).transpose();
        PearsonsCorrelation pearsonsCorrelation = new PearsonsCorrelation(dataMatrix);
        correlationMatrix = pearsonsCorrelation.getCorrelationMatrix();
        correlationResultList = new ArrayList<>();
//        for (int i = 0; i < correlationMatrix.getColumnDimension(); i++) {
//            for (int j = i+1; j < correlationMatrix.getRowDimension(); j++) {
//                log.debug("correlation[{}][{}]: {}", i, j, correlationMatrix.getEntry(i, j));
//            }
//        }
        RealMatrix correlationPValues = pearsonsCorrelation.getCorrelationPValues();
//        for (int i = 0; i < correlationPValues.getColumnDimension(); i++) {
//            for (int j = i+1; j < correlationPValues.getRowDimension(); j++) {
//                log.debug("PValues[{}][{}]: {}", i, j, correlationPValues.getEntry(i, j));
//            }
//        }


        // 2. 计算每个相关系数的置信区间
        int n = data[0].length; // 样本数量
        log.debug("n: "+n);
        double confidenceLevel = 0.95; // 置信水平

        // 创建t分布对象
//        TDistribution tDistribution = new TDistribution(n - 2);
//        double criticalValue = 1.96;
//        double criticalValue = tDistribution.inverseCumulativeProbability(1 - (1 - confidenceLevel)/2);
        NormalDistribution normalDistribution = new NormalDistribution(0, 1);
        double criticalValue = normalDistribution.inverseCumulativeProbability(1 - (1 - confidenceLevel) / 2);
        log.debug("criticalValue: {}", criticalValue);
        // 遍历相关系数矩阵
        for (int i = 0; i <  correlationMatrix.getColumnDimension(); i++) {
            for (int j = i+1; j < correlationMatrix.getRowDimension(); j++) {
//                if (i != j) {
                    double r = correlationMatrix.getEntry(i,j);
                    double pValue = correlationPValues.getEntry(i,j);
                    // 计算置信区间
                    double[] ci = calculateConfidenceInterval(r, n, criticalValue);
                    log.debug("Correlation between {} and {}: {}, 95%% CI: [{}, {}]\n", i, j, r, ci[0], ci[1]);
//                }
                Map<String,String> correlationResult = new LinkedHashMap<>();
                correlationResult.put("column 1", String.valueOf(i));
                correlationResult.put("column 2", String.valueOf(j));
                correlationResult.put("N", String.valueOf(n));
                correlationResult.put("r", String.format("%.4f",r));
                correlationResult.put("ci", "["+String.format("%.4f",ci[0])+","+String.format("%.4f",ci[1])+"]");
                correlationResult.put("pValue", String.format("%.4f",pValue));
                correlationResultList.add(correlationResult);

            }
        }

    }


    // 计算相关系数的置信区间
    private static double[] calculateConfidenceInterval(double r, int n, double criticalValue) {
        // Fisher z变换
        double z = 0.5 * Math.log((1 + r) / (1 - r));

        // 计算标准误差
        double se = 1.0 / Math.sqrt(n - 3);

        // 计算z的置信区间
        double lowerZ = z - criticalValue * se;
        double upperZ = z + criticalValue * se;

        // 逆Fisher变换得到相关系数的置信区间
        double lowerR = (Math.exp(2 * lowerZ) - 1) / (Math.exp(2 * lowerZ) + 1);
        double upperR = (Math.exp(2 * upperZ) - 1) / (Math.exp(2 * upperZ) + 1);

        return new double[]{lowerR, upperR};
    }
    // 计算相关系数的置信区间
    private static double[] calculateConfidenceIntervalOfSpearman(double r, int n, double criticalValue) {
        // Fisher z变换
        double z = 0.5 * Math.log((1 + r) / (1 - r));

        // 计算标准误差
        double se = Math.sqrt(((1.0 + 0.5* r * r) / (n - 3)));

        // 计算z的置信区间
        double lowerZ = z - criticalValue * se;
        double upperZ = z + criticalValue * se;

        // 逆Fisher变换得到相关系数的置信区间
        double lowerR = (Math.exp(2 * lowerZ) - 1) / (Math.exp(2 * lowerZ) + 1);
        double upperR = (Math.exp(2 * upperZ) - 1) / (Math.exp(2 * upperZ) + 1);

        return new double[]{lowerR, upperR};
    }

    public Figure getScatterPlot(double[] x, double[] y, String xTitle, String yTitle ){
        return ScatterPlot.create("", xTitle, x, yTitle, y);
    }


    public Figure[][]  getScatterPlots(double[][] data, String []columnName){

        RealMatrix dataMatrix = MatrixUtils.createRealMatrix(data).transpose();
        Figure[][] figureMatrix = new Figure[dataMatrix.getColumnDimension()][dataMatrix.getColumnDimension()];
        for (int i=0;i<dataMatrix.getColumnDimension();i++){
            for (int j=0;j<dataMatrix.getColumnDimension();j++){
                String xTitle = "第"+ i +"列",yTitle =  "第"+ j +"列";
                if (columnName != null&&columnName.length==dataMatrix.getColumnDimension()){
                    xTitle = columnName[i];
                    yTitle = columnName[j];
                }
                Figure scatterPlot = ScatterPlot.create("", xTitle, dataMatrix.getColumn(i), yTitle, dataMatrix.getColumn(j));
                figureMatrix[i][j] = scatterPlot;
//                Plot.show(scatterPlot);
            }
        }
        return figureMatrix;
    }
}
